import warnings
warnings.filterwarnings("ignore", category=Warning)
from torch_geometric.transforms import NormalizeFeatures
from torch_geometric.datasets import Reddit2, Yelp, AttributedGraphDataset

# dataset = Reddit2(
#     root='data/Reddit2',  # 自动识别子目录 reddit2/
#     transform=None,
#     pre_transform=None,
#     force_reload=False
# )

# dataset = Yelp(
#     root='./data/Yelp',
#     transform=NormalizeFeatures()  # 自动处理0维特征
# )

# # 加载 TWeibo 数据集
# dataset = AttributedGraphDataset(
#     root='data/TWeibo',        # 数据集存储路径
#     name='TWeibo',             # 指定数据集名称（必须与文档中STATS表对应）
#     transform=NormalizeFeatures(),  # 特征标准化（适用于157维特征）
#     pre_transform=None,        # 预处理操作（可选）
#     force_reload=False         # 是否强制重新处理
# )

# 加载 MAG 数据集
dataset = AttributedGraphDataset(
    root='data/MAG',        # 数据集存储路径
    name='MAG',             # 指定数据集名称（必须与文档中STATS表对应）
    transform=NormalizeFeatures(),  # 特征标准化（适用于157维特征）
    pre_transform=None,        # 预处理操作（可选）
    force_reload=False         # 是否强制重新处理
)

data = dataset[0]  # 获取第一个图数据

# 核心属性检查
print("节点特征矩阵形状:", data.x.shape)        # (716847, 300)
print("边索引形状:", data.edge_index.shape)     # (2, 13954819)
print("标签分布:", dataset.y.unique(return_counts=True))  # 100类分布统计